# coding: utf-8
import numpy as np
import shutil
import os
import sys
import time
if sys.version_info[0] == 2:
    import xml.etree.cElementTree as ET
else:
    import xml.etree.ElementTree as ET
from common import *

from utils.Logger import Logger

if not os.path.exists("./log/"):
    os.mkdir("./log/")
log = Logger("./log/{}_{}.log".format(__file__.split('/')[-1],
                                             time.strftime("%Y%m%d-%H%M%S"), time.localtime), level='debug').logger


source_dir = "~/dataset/VOCdevkit/VOC2012"
source_dir = os.path.expanduser(source_dir)
output_dir = DATASET_ROOT_DIR + "/voc2012"

# mkdir
mk_common_dir(output_dir)

# analysis source
fp = os.path.join(source_dir, "ImageSets/Main/trainval.txt")
ids = []
with open(fp) as f:
    ids = f.readlines()

#convert
global_idx = 0
for i in ids:
    xml_path = os.path.join(source_dir, "Annotations/{}.xml".format(i.strip()))
    # read xml
    target = ET.parse(xml_path).getroot()

    # image size
    c = 0
    h = 0
    w = 0
    for image_size in target.iter('size'):
        w = int(image_size.find('width').text)
        h = int(image_size.find('height').text)
        c = int(image_size.find('depth').text)
        break

    object_names = []
    object_bboxes = []
    for obj in target.iter('object'):
        name = obj.find('name').text.lower().strip()
        if name not in voc_label:
            continue

        bbox = obj.find('bndbox')
        pts = ['xmin', 'ymin', 'xmax', 'ymax']
        bndbox = []
        for pt in pts:
            cur_pt = int(bbox.find(pt).text)
            bndbox.append(cur_pt)
        object_names += [name]
        object_bboxes += [bndbox]
    # write xml
    if len(object_bboxes) > 0:
        log.info("{} have object: {}".format(i.strip(), object_names))
        image_name = "voc2012_{}".format(global_idx)
        global_idx += 1

        #
        write_xml(os.path.join(output_dir, "Annotations"), image_name, "voc2012", [w, h, c], object_bboxes, object_names)

        # copy image
        image_path = os.path.join(output_dir, "JPEGImages/{}.jpg".format(image_name))
        shutil.copy(os.path.join(source_dir, "JPEGImages/{}.jpg".format(i.strip())), image_path)

log.info("process done, whole number is: {}".format(global_idx))
